`timescale 1ns / 1ps

module MUL_MOD(
    input  wire [11:0] a , b,
    input  wire clk,en,rst_n,
    output reg [11:0] product,

    input   mul_mode, // mul mode select, 0-as MUL MOD, 1 as-compress/decompress
    input   Xpress_mode, // 0-compress, 1-decompress
    input   d_mode, // d_mode, 0 = du, 1 = dv
    input   k_mode, // k_mode, 0 = kyber512/768, 1 = kyber1024
                    // for d_mode and k_mode select, the state is below 
                    //        0         0            du10,  du=10 kyber512/768
                    //        1         0            dv4,   dv=4  kyber512/768
                    //        0         1            du11,  du=11 kyber1024
                    //        1         1            dv5,   dv=5  kyber1024

    input  [11:0]   din_compress,
    output [11:0]   do_compress,

    input  [11:0]   din_decompress,
    output [11:0]   do_decompress

);

//////////////////////////////////////////////////////////////////
// 第一级
// reg [23:0] mul_reg; //  a x b
// always @(posedge clk or negedge rst_n) begin
//     if(!rst_n) mul_reg<=0;
//     else if(en)   mul_reg<=a*b;
// end

wire [23:0] mul_reg; //  a x b


wire [21:0] Xpress_A;
wire [11:0] Xpress_B;
wire [33:0] Xpress_P;

wire [21:0] Compress_A;
wire [11:0] Compress_B;
wire [33:0] Compress_P;

wire [21:0] mul_A_in;
wire [11:0] mul_B_in;
wire [33:0] mul_P_out;

wire [21:0] Decompress_A;
wire [11:0] Decompress_B;
wire [33:0] Decompress_P;




///////////////TEST////////////////////////////////////
reg [21:0] Xpress_A_buf;
reg [11:0] Xpress_B_buf;
always@(posedge clk or negedge rst_n) begin
    if(!rst_n) begin
        Xpress_A_buf <= 22'd0;
        Xpress_B_buf <= 12'd0;
    end
    else begin
        Xpress_A_buf <= Xpress_A;
        Xpress_B_buf <= Xpress_B;
    end
end

// reg [21:0] mul_A_in;
// reg [11:0] mul_B_in;
// always@(posedge clk or negedge rst_n) begin
//     if(!rst_n) begin
//         mul_A_in <= 22'd0;
//         mul_B_in <= 12'd0;
//     end
//     else begin
//         mul_A_in <= mul_mode ? Xpress_A : {10'd0,a};
//         mul_B_in <= mul_mode ? Xpress_B : b;
//     end
// end

assign mul_A_in = mul_mode ? Xpress_A_buf : {10'd0,a};
assign mul_B_in = mul_mode ? Xpress_B_buf : b;
///////////////TEST////////////////////////////////////



// assign mul_A_in = mul_mode ? Xpress_A : {10'd0,a};
// assign mul_B_in = mul_mode ? Xpress_B : b;



assign mul_reg = mul_P_out[23:0];
assign Xpress_P = mul_P_out;

assign Xpress_A = Xpress_mode ? Decompress_A : Compress_A;
assign Xpress_B = Xpress_mode ? Decompress_B : Compress_B;

assign Compress_P = Xpress_P;
assign Decompress_P = Xpress_P;

MUL22x12 MUL22x12_u0 (
    .clk(clk),
    .A(mul_A_in),
    .B(mul_B_in),
    .P(mul_P_out)
    );

// DSP_22x12 u_DSP_22x12 (
//   .CLK(clk),  // input wire CLK
//   .A(mul_A_in),      // input wire [21 : 0] A
//   .B(mul_B_in),      // input wire [11 : 0] B
//   .P(mul_P_out)      // output wire [33 : 0] P
// );

//////////////////////////////////////////////////////////////////


//////////////////////////////////////////////////////////////////
Kyber_Compress Kyber_Compress_u0(

    .d_mode(d_mode), // d_mode, 0 = du, 1 = dv
    .k_mode(k_mode), // k_mode, 0 = kyber512/768, 1 = kyber1024
                    // for d_mode and k_mode select, the state is below 
                    //        0         0            du10,  du=10 kyber512/768
                    //        1         0            dv4,   dv=4  kyber512/768
                    //        0         1            du11,  du=11 kyber1024
                    //        1         1            dv5,   dv=5  kyber1024
    .Mul_A(Compress_A),
    .Mul_B(Compress_B),
    .Mul_P(Compress_P),

    .data_in(din_compress),
    .data_out(do_compress)
    );

//////////////////////////////////////////////////////////////////


//////////////////////////////////////////////////////////////////
Kyber_Decompress Kyber_Decompress_u0(

    .d_mode(d_mode), // d_mode, 0 = du, 1 = dv
    .k_mode(k_mode), // k_mode, 0 = kyber512/768, 1 = kyber1024
                    // for d_mode and k_mode select, the state is below 
                    //        0         0            du10,  du=10 kyber512/768
                    //        1         0            dv4,   dv=4  kyber512/768
                    //        0         1            du11,  du=11 kyber1024
                    //        1         1            dv5,   dv=5  kyber1024
    .Mul_A(Decompress_A),
    .Mul_B(Decompress_B),
    .Mul_P(Decompress_P),

    .data_in(din_decompress),
    .data_out(do_decompress)
    );
//////////////////////////////////////////////////////////////////



wire [3:0] mul_h , mul_m, mul_l;
wire [11:0] mul_q;
assign {mul_h, mul_m, mul_l, mul_q} = mul_reg;


//第二级
wire [11:0] lut_h, lut_m,lut_l;
wire [12:0] add13_sum0,add13_sum1;
wire [13:0] add14_sum;
reg [13:0] add14_sum_reg;
ROM_H rom0 (
.a (mul_h),
.spo (lut_h)
);

ROM_M rom1 (
.a (mul_m),
.spo (lut_m)
);

ROM_L rom2 (
.a (mul_l),
.spo (lut_l)
);

assign add13_sum0 = lut_h + lut_m;
assign add13_sum1 = lut_l + mul_q;
assign add14_sum   = add13_sum0 + add13_sum1;

always @( posedge clk or negedge rst_n) begin 
    if(!rst_n) add14_sum_reg<=0;
    else if(en) add14_sum_reg<=add14_sum;
end


//第三级
reg [11:0] ma0_j;
wire flag;
wire [11:0] sub,ma0_k,ma0_s;
always@(add14_sum_reg)
  begin
    case(add14_sum_reg[13:12])
	    2'b00: ma0_j = 12'd0;
	    2'b01: ma0_j = 12'd767;
	    2'b10: ma0_j = 12'd1534;
	    2'b11: ma0_j = 12'd2301;		
	endcase
  end

assign {flag,sub} = add14_sum_reg[11:0] - 12'd3329;
assign ma0_k = flag? add14_sum_reg[11:0] : sub;

MA ma0 (
.MA_a(ma0_j),
.MA_b(ma0_k),
.MA_s(ma0_s)
);

//第四级

always @( posedge clk or negedge rst_n) begin
    if(!rst_n) product<=0;
    else if(en ) product<=ma0_s;
end
endmodule



module MUL_MOD_tb;
reg clk,rst_n,en;
reg [11:0] a, b;
reg [22:0] check1,check2,check3;
wire  [11:0] product;


MUL_MOD themul_mod(.a(a),.b(b),.clk(clk),.en(en),.rst_n(rst_n),.product(product));


initial begin
    clk=0;
    rst_n=1;
    en=1;
    a=0;
    b=0;
    #10 rst_n=0;
end
always #5 clk=~clk;

always@(posedge clk)
begin
    a<=({$random}%3329);
    b<=({$random}%3329);
    check1<=a*b%3329;
    check2<=check1;
    check3<=check2;
end

endmodule